from copy import deepcopy

import trimesh
import networkx as nx
import numpy as np

def build_graph(mesh):
    edges = mesh.edges_unique
    length = mesh.edges_unique_length
    g = nx.Graph()
    for edge, L in zip(edges, length):
        g.add_edge(*edge, length=L)
    return g

def fps(points, npoint, mesh, use_geodesic=False):
    """
    Input:
        mesh: input mesh
        graph: graph for mesh
        npoint: target point number to sample
    Return:
        centroids: sampled pointcloud index, [npoint]
    """
    print(npoint)
    assert use_geodesic == False
    if use_geodesic:
        graph = build_graph(mesh)
        N, C = mesh.vertices.shape
    else:
        N, C = points.shape
    centroids = np.zeros(npoint, dtype=np.int)
    distance = np.ones(N) * 1e10
    farthest = np.random.randint(0, N)
    for i in range(npoint):
        centroids[i] = farthest
        # centroid = mesh.vertices[farthest, :].reshape(1, 3)
        centroid = points[farthest, :].reshape(1, 3)
        if not use_geodesic:
            # dist = np.sum((mesh.vertices - centroid) ** 2, -1)
            dist = np.sum((points - centroid) ** 2, -1)
            mask = dist < distance
            distance[mask] = dist[mask]
        else:
            dist = nx.shortest_path_length(graph, source=farthest, weight='length')
            # dist = length_geodesic
            for idx in range(0, N):
                if dist[idx] < distance[idx]:
                    distance[idx] = dist[idx]

        farthest = np.argmax(distance, -1)
    return centroids

def sample_meshpoints(vis, numpoints, mesh_file='new_mesh.stl'):
    mesh = trimesh.load_mesh(mesh_file)
    sampled_points, tri_index = mesh.sample(1000, return_index=True)
    sample_point_index = fps(sampled_points, numpoints, mesh, use_geodesic=False)
    points = sampled_points[sample_point_index]

    # normals = mesh.vertex_normals[sample_point_index]
    # result = np.concatenate((points, -normals), axis=1)

    if vis:
        pc = trimesh.PointCloud(points, colors=[255, 0, 0])
        ray_origins = points
        # ray_directions = normals
        # vis_path = np.hstack((ray_origins, ray_origins + ray_directions / 50)).reshape(-1, 2, 3)
        # ray_visualize = trimesh.load_path(vis_path)
        scene = trimesh.Scene()
        scene.add_geometry(pc)
        scene.add_geometry(mesh)
        # scene.add_geometry(ray_visualize)
        scene.show()
    return None


if __name__ == "__main__":
    import os

    root_dir = '/home/v-wewei/code/new_mesh'
    for root, dirs, files in os.walk(root_dir):
        for file_name in files:
            if not file_name.startswith('large_mug'):
                continue
            file_path = os.path.join(root, file_name)

            mesh = trimesh.load_mesh(file_path)

            mesh.vertices *= 2
            mesh.export(file_name)
            # print(ori_mesh_box_axis,mesh_box_axis)
            # exit()
    # sections = 50
    # annular = trimesh.creation.annulus(0.03, 0.035, height=0.015, sections=sections, transform=None, segment=None)
    # mesh = trimesh.load_mesh('new_mesh.stl')
    # new_mesh = annular.slice_plane([0, 0, 0], [0, 1, 0])
    # res = sample_meshpoints(vis=True,numpoints=500)
    # scene = trimesh.Scene()
    # T = trimesh.transformations.translation_matrix([0, 0.0, 0.000])
    # annular = trimesh.creation.annulus(0.028, 0.03, height=0.055, sections=sections, transform=T, segment=None)
    # T1 = trimesh.transformations.translation_matrix([0, 0.0, -0.03])
    # plane = trimesh.creation.cylinder(0.03, height=0.005, sections=sections, segment=None, transform=None)
    # mug = trimesh.load_mesh("/home/v-wewei/finish_stl/025_mug.stl")
    # mug.visual.face_colors = [0, 255, 255]
    # plane.visual.face_colors = [0, 255, 255]
    # scene.add_geometry(new_mesh)
    # scene.add_geometry(mesh)
    # scene.add_geometry(plane)
    # mesh_combine =annular+plane
    # mesh_combine.export('/home/lifuyu/cup.stl')

    # T3 = trimesh.transformations.translation_matrix([0, -0.01, 0.0])
    # box = trimesh.creation.box([0.06, 0.02, 0.005], transform=T3)
    # T4 = trimesh.transformations.translation_matrix([-0.0325, -0.05, 0.0])
    # box1 = trimesh.creation.box([0.005, 0.1, 0.015], transform=T4)
    # box2 = deepcopy(box1)
    # T5 = trimesh.transformations.translation_matrix([0.065, 0.0, 0.0])
    # box2.apply_transform(T5)
    # scene.add_geometry(new_mesh)
    # scene.add_geometry(box)
    # scene.add_geometry(box1)
    # scene.add_geometry(box2)
    # scene.show()
    # mesh_combine = new_mesh + mesh + box1 + box2
    # mesh_combine = new_mesh+box
    # mesh_combine.export('/home/lifuyu/niezi.stl')
